from utils.global_utils import *


def white_bg(h: int, w: int) -> np.ndarray:
    """
    生成全白图像。
    Args:
        h (int): 图片的高度H
        w (int): 图片的高度W
    Returns:
        img (np.ndarray): 全白图像
    """
    return np.full((h, w, 3), 255, dtype = "uint8")


def black_bg(h: int, w: int) -> np.ndarray:
    """
    生成全黑图像。
    Args:
        h (int): 图片的高度H
        w (int): 图片的高度W
    Returns:
        img (np.ndarray): 全黑图像
    """
    return np.zeros((h, w, 3), dtype = "uint8")


def bgr2hex(bgr: np.ndarray) -> str:
    """
    将BGR颜色转成十六进制颜色字符串。
    Args:
        bgr (np.ndarray): BGR颜色向量
    Returns:
        hex (str): 十六进制颜色字符串
    """
    bgr = bgr.flatten()
    b, g, r = bgr[0], bgr[1], bgr[2]
    res = (r << 16) + (g << 8) + b
    res = hex(res).replace("0x", "")
    res = "0" * (6 - len(res)) + res
    return "#" + res


def hex2bgr(hex: str) -> np.ndarray:
    """
    将十六进制颜色字符串转成BGR颜色。
    Args:
        hex (str): 十六进制颜色字符串
    Returns:
        bgr (np.ndarray): BGR颜色向量
    """
    hex = int(hex.replace("0x", "").replace("#", ""), base = 16)
    bgr = np.array([[[hex & 0xff, (hex >> 8) & 0xff, (hex >> 16) & 0xff]]], dtype = "uint8")
    return bgr


def bgr2hsv(bgr: np.ndarray) -> np.ndarray:
    """
    将BGR颜色转成HSV颜色。
    Args:
        bgr (np.ndarray): BGR颜色向量
    Returns:
        hsv (np.ndarray): HSV颜色向量
    """
    hsv = cv2.cvtColor(bgr.astype("uint8"), cv2.COLOR_BGR2HSV)
    return hsv


def hsv2bgr(hsv: np.ndarray) -> np.ndarray:
    """
    将HSV颜色转成BGR颜色。
    Args:
        hsv (np.ndarray): HSV颜色向量
    Returns:
        bgr (np.ndarray): BGR颜色向量
    """
    bgr = cv2.cvtColor(hsv.astype("uint8"), cv2.COLOR_HSV2BGR)
    return bgr


def hsv2hex(hsv: np.ndarray) -> str:
    """
    将HSV颜色转成十六进制颜色字符串。
    Args:
        hsv (np.ndarray): HSV颜色向量
    Returns:
        hex (str): 十六进制颜色字符串
    """
    return bgr2hex(hsv2bgr(hsv))


def hex2hsv(hex: str) -> np.ndarray:
    """
    将十六进制颜色字符串转成HSV颜色。
    Args:
        hex (str): 十六进制颜色字符串
    Returns:
        hsv (np.ndarray): HSV颜色向量
    """
    return bgr2hsv(hex2bgr(hex))


def add_events(img: np.ndarray, events: np.ndarray) -> np.ndarray:
    """
    添加事件。
    Args:
        img (np.ndarray): 空白图片，形状为[H, W, 3]
        events (np.ndarray): 事件，形状为[N, 4]
    Returns:
        rendered_img (np.ndarray): 渲染后的图片，形状为[H, W, 3]
    """
    pos_color = np.array([0, 0, 255], dtype = np.float32)
    neg_color = np.array([255, 0, 0], dtype = np.float32)

    p = events[:, 1]
    pos = events[p > 0.5, 2:]
    neg = events[p <= 0.5, 2:]
    idx_pos, val_pos = np.unique(pos, axis = 0, return_counts = True)
    idx_neg, val_neg = np.unique(neg, axis = 0, return_counts = True)
    idx_pos = idx_pos.astype(np.int64)
    idx_neg = idx_neg.astype(np.int64)
    val_pos = val_pos.astype(np.float32)
    val_neg = val_neg.astype(np.float32)
    val_max = max(np.max(val_pos) if len(val_pos) else 0, np.max(val_neg) if len(val_neg) else 0, 1)
    val_pos = np.stack([val_pos] * 3, axis = 1) / val_max
    val_neg = np.stack([val_neg] * 3, axis = 1) / val_max

    pos_frame = img[idx_pos[:, 0], idx_pos[:, 1]].astype(np.float32)
    pos_frame = pos_color * val_pos + pos_frame * (1.0 - val_pos)
    img[idx_pos[:, 0], idx_pos[:, 1]] = pos_frame.astype(np.uint8)
    neg_frame = img[idx_neg[:, 0], idx_neg[:, 1]].astype(np.float32)
    neg_frame = neg_color * val_neg + neg_frame * (1.0 - val_neg)
    img[idx_neg[:, 0], idx_neg[:, 1]] = neg_frame.astype(np.uint8)

    return img